
#include "pcg_def.h"
#include <slave.h>

#include "my_slave.h"
#include <crts.h>

#include <math.h>

#include <simd.h>

#define MIN(a, b) ((b) < (a) ? (b) : (a))

__thread_local crts_rply_t DMARply = 0;
__thread_local unsigned int DMARplyCount = 0;

typedef struct {
    // 总行数
    int total_rows;
    // 矩阵切片的起始行索引
    int row_base;
    // 矩阵切片的行数
    int rows;
    // 上下三角矩阵中每一行的偏移量（长为 rows + 1）
    uint16_t* row_off1;
    uint16_t* row_off2;
    // 上下三角矩阵中每一个元素的列索引
    int* cols1;
    int* cols2;
    // 上下三角矩阵的元素值
    double* data1;
    double* data2;
    // 对角线元素
    double* diag;
    double* source;
} MatrixSlice;

#define VEC_BUF_SIZE 88500

// 当前从核保存的 CSR 矩阵切片
__thread_local MatrixSlice matrix_slice;
__thread_local_share double x[VEC_BUF_SIZE];

__thread_local_share double Ax[VEC_BUF_SIZE];
__thread_local_share double r[VEC_BUF_SIZE];
__thread_local_share double z[VEC_BUF_SIZE];
__thread_local_share double p[VEC_BUF_SIZE];

static int* lower_bound(int* first, int* last, int value) {
    int* it;
    int count = last - first, step;

    while (count > 0) {
        it = first;
        step = count / 2;
        it += step;

        if (*it < value) {
            first = ++it;
            count -= step + 1;
        } else
            count = step;
    }

    return first;
}

// 从核数组的简单分段归约实现
__thread_local uint8_t* buf_ptr;
__thread_local athread_rply_t rma_rply = 0;
__thread_local int rma_rply_count = 0;
void simple_segmented_reduce(int remote, int buf_size) {
#define CHUNK_SIZE 2048
    uint8_t buf1[CHUNK_SIZE], buf2[CHUNK_SIZE];
    uint8_t *buf_rma = buf1, *buf_process = buf2;
    uint8_t* buf_result = buf_ptr;
    uint8_t* buf_remote;
    athread_rma_get(&buf_remote, sizeof(buf_remote), remote, &buf_ptr, &rma_rply);

    // 第一次加载
    athread_rma_iget(
        buf_rma,
        &rma_rply,
        MIN(CHUNK_SIZE, buf_size) * sizeof(uint8_t),
        remote,
        buf_remote,
        &rma_rply
    );
    ++rma_rply_count;

    int i = CHUNK_SIZE;
    for (; i < buf_size; i += CHUNK_SIZE) {
        athread_rma_wait_value(&rma_rply, rma_rply_count);
        // 交换结果指针
        {
            uint8_t* temp = buf_rma;
            buf_rma = buf_process;
            buf_process = temp;
        }
        // 开始下一段 RMA
        athread_rma_iget(
            buf_rma,
            &rma_rply,
            MIN(CHUNK_SIZE, buf_size - i) * sizeof(uint8_t),
            remote,
            buf_remote + i,
            &rma_rply
        );
        ++rma_rply_count;
        // 处理当前段
        for (int j = 0; j != CHUNK_SIZE; ++j) {
            buf_result[i + j - CHUNK_SIZE] += buf_process[j];
        }
    }

    // 处理最后一段
    athread_rma_wait_value(&rma_rply, rma_rply_count);
    for (int off = i - CHUNK_SIZE, j = 0; j < buf_size - off; ++j) {
        buf_result[j + off] += buf_rma[j];
    }
#undef CHUNK_SIZE
}

__thread_local_share int thread_data_off[65];

// 载入输入参数到 LDM 中 - 用于和 ldu_to_csr 并行执行
void load_input_param1(LoadInputParam1* _param) {
    LoadInputParam1 param;
    athread_dma_get(&param, _param, sizeof(param));

    int const rest = param.cells % 64;
    int const len = _PEN < rest ? param.cells / 64 + 1 : param.cells / 64;
    int const rows_beg = _PEN < rest ? _PEN * len : _PEN * len + rest;
    int const rows_end = rows_beg + len;

    matrix_slice.total_rows = param.cells;
    matrix_slice.row_base = rows_beg;
    matrix_slice.rows = len;

    /* 加载向量 */
    athread_memcpy_sldm(x, param.x, matrix_slice.total_rows * sizeof(double), MEM_TO_LDM);

    matrix_slice.diag = ldm_malloc(len * sizeof(double));
    matrix_slice.source = ldm_malloc(len * sizeof(double));
    athread_dma_iget(matrix_slice.diag, param.diag + rows_beg, len * sizeof(double), &DMARply);
    athread_dma_iget(matrix_slice.source, param.source + rows_beg, len * sizeof(double), &DMARply);

    // 加载向量时，计算并加载 LDU 矩阵

    // 计算当前线程处理的元素范围
    if (_PEN == 0) {
        thread_data_off[0] = 0;
        thread_data_off[64] = param.faces;
    } else {
        thread_data_off[_PEN] =
            lower_bound(param.lPtr, param.lPtr + param.faces, rows_beg) - param.lPtr;
    }
    athread_ssync_array();

    // 载入元素
    int const data_beg = thread_data_off[_PEN];
    int const data_end = thread_data_off[_PEN + 1];
    int const elem_count = data_end - data_beg;
    matrix_slice.cols1 = ldm_malloc(elem_count * sizeof(int));
    matrix_slice.data1 = ldm_malloc(elem_count * sizeof(double));
    athread_dma_iget(
        matrix_slice.cols1,
        param.uPtr + data_beg,
        elem_count * sizeof(int),
        &DMARply
    );
    athread_dma_iget(
        matrix_slice.data1,
        param.data + data_beg,
        elem_count * sizeof(double),
        &DMARply
    );
    DMARplyCount += 4;

    // 等待 DMA 的同时构造 row_off1
    matrix_slice.row_off1 = ldm_malloc((rows_end - rows_beg + 1) * sizeof(uint16_t));
    matrix_slice.row_off1[0] = 0;
    matrix_slice.row_off1[rows_end - rows_beg] = elem_count;

    int* const thread_lptr_beg = param.lPtr + data_beg;
    int* const thread_lptr_end = param.lPtr + data_end;
    for (int i = 1; i != rows_end - rows_beg; ++i) {
        // 二分搜索的范围依赖于上一次搜索结果
        // 这会减少搜索范围，但使得循环间产生了数据依赖
        // TODO: 性能测试
        matrix_slice.row_off1[i] = lower_bound(
                                       thread_lptr_beg + matrix_slice.row_off1[i - 1],
                                       thread_lptr_end,
                                       i + rows_beg
                                   )
            - thread_lptr_beg;
    }

    athread_dma_wait_value(&DMARply, DMARplyCount);
}

void load_input_param2(struct CsrMatrix* _param) {
    struct CsrMatrix param;
    athread_dma_get(&param, _param, sizeof(param));

    int const rows_beg = matrix_slice.row_base;
    int const rows_end = matrix_slice.row_base + matrix_slice.rows;

    // 加载 row_off
    matrix_slice.row_off2 = ldm_malloc((matrix_slice.rows + 1) * sizeof(uint16_t));
    int* temp_buf = ldm_malloc((matrix_slice.rows + 1) * sizeof(int));
    athread_dma_get(temp_buf, param.row_off + rows_beg, (matrix_slice.rows + 1) * sizeof(int));

    int const row_off_base = temp_buf[0];
    for (int i = 0; i != matrix_slice.rows + 1; ++i) {
        matrix_slice.row_off2[i] = (uint16_t) (temp_buf[i] - row_off_base);
    }

    ldm_free(temp_buf, (matrix_slice.rows + 1) * sizeof(int));

    // 加载 data 和 cols
    int const elem_count = matrix_slice.row_off2[matrix_slice.rows];
    matrix_slice.cols2 = ldm_malloc(elem_count * sizeof(int));
    matrix_slice.data2 = ldm_malloc(elem_count * sizeof(double));

    athread_dma_get(matrix_slice.cols2, param.cols + row_off_base, elem_count * sizeof(int));
    athread_dma_get(matrix_slice.data2, param.data + row_off_base, elem_count * sizeof(double));
}

void free_csr_matrix_in_ldm() {
    int const elem_count1 = matrix_slice.row_off1[matrix_slice.rows];
    int const elem_count2 = matrix_slice.row_off2[matrix_slice.rows];
    ldm_free(matrix_slice.data2, elem_count2 * sizeof(double));
    ldm_free(matrix_slice.cols2, elem_count2 * sizeof(int));
    ldm_free(matrix_slice.row_off2, (matrix_slice.rows + 1) * sizeof(uint16_t));

    ldm_free(matrix_slice.row_off1, (matrix_slice.rows + 1) * sizeof(uint16_t));
    ldm_free(matrix_slice.data1, elem_count1 * sizeof(double));
    ldm_free(matrix_slice.cols1, elem_count1 * sizeof(int));

    ldm_free(matrix_slice.source, matrix_slice.rows * sizeof(double));
    ldm_free(matrix_slice.diag, matrix_slice.rows * sizeof(double));
}

void csr_spmv(double* input_vec, double* output_vec) {
    int const row_base = matrix_slice.row_base;
    for (int i = 0; i != matrix_slice.rows; ++i) {
        // 对角线
        double temp = matrix_slice.diag[i] * input_vec[i + row_base];
        // 三角 1
        for (int j = matrix_slice.row_off1[i], end = matrix_slice.row_off1[i + 1]; j != end; ++j) {
            temp += input_vec[matrix_slice.cols1[j]] * matrix_slice.data1[j];
        }
        // 三角 2
        for (int j = matrix_slice.row_off2[i], end = matrix_slice.row_off2[i + 1]; j != end; ++j) {
            temp += input_vec[matrix_slice.cols2[j]] * matrix_slice.data2[j];
        }
        output_vec[i + row_base] = temp;
    }
}

void csr_precondition_spmv(double* input_vec, double* output_vec) {
    int const row_base = matrix_slice.row_base;
    for (int i = 0; i != matrix_slice.rows; ++i) {
        // 对角线 - 全 0
        double temp = 0;
        // 三角 1
        for (int j = matrix_slice.row_off1[i], end = matrix_slice.row_off1[i + 1]; j != end; ++j) {
            temp += input_vec[matrix_slice.cols1[j]] * matrix_slice.data1[j];
        }
        // 三角 2
        for (int j = matrix_slice.row_off2[i], end = matrix_slice.row_off2[i + 1]; j != end; ++j) {
            temp += input_vec[matrix_slice.cols2[j]] * matrix_slice.data2[j];
        }
        // 不同于 csr_spmv，这里的下标是 i 而不是 i + row_base，因为输出向量是局部的
        output_vec[i] = temp;
    }
}

void pcg_precondition_csr(double* rAPtr, double* wAPtr) {
    int const row_base = matrix_slice.row_base;

    for (int i = row_base, end = matrix_slice.rows + row_base; i != end; ++i) {
        wAPtr[i] = rAPtr[i] / matrix_slice.diag[i - row_base];
    }

    athread_ssync_array();

    double gAPtr[1500];
    csr_precondition_spmv(wAPtr, gAPtr);
    athread_ssync_array();

    for (int i = row_base, end = matrix_slice.rows + row_base; i != end; ++i) {
        wAPtr[i] = (rAPtr[i] - gAPtr[i - row_base]) / matrix_slice.diag[i - row_base];
    }
}

double pcg_gsumProd_imbalanced(double* z, double* r) {
    int const row_base = matrix_slice.row_base;
    double local_result = 0.0, result;
    {
        doublev8 simd_result = .0, simd_z, simd_r;
        int len = matrix_slice.rows;
        for (int i = 0; i < len / 8; i++) {
            simd_loadu(simd_z, &(z[i * 8 + row_base]));
            simd_loadu(simd_r, &(r[i * 8 + row_base]));
            simd_result += simd_z * simd_r;
        }
        for (int i = (len / 8) * 8; i < len; i++) {
            local_result += z[i + row_base] * r[i + row_base];
        }
        local_result += simd_reduc_plusd(simd_result);
    }
    double reduce_buf[64];
    athread_ssync_array();
    athread_redurt(&local_result, &result, 1, athread_double, OP_add, reduce_buf, 64);

    return result;
}

void run(RunParam* _param) {
    RunParam param;
    athread_dma_get(&param, _param, sizeof(param));

    load_input_param2(param.csr_matrix);

    double residual = 0.0;
    int iter = 0;
    double sumprod_old = 0.0, sumprod = 0.0;
    double alpha = 0.0, beta = 0.0;

    int const row_base = matrix_slice.row_base;

    // AX = A * X
    csr_spmv(x, Ax);

    // r = b - A * x
    {
        double _local_residual = 0.0;

        for (int i = row_base, end = matrix_slice.rows + row_base; i != end; ++i) {
            r[i] = matrix_slice.source[i - row_base] - Ax[i];
        }

        {
            int len = matrix_slice.rows;
            doublev8 simd_abs_r;
            for (int i = 0; i < len / 8; i++) {
                simd_abs_r = simd_set_doublev8(
                    fabs(r[i * 8 + row_base]), 
                    fabs(r[i * 8 + row_base + 1]), 
                    fabs(r[i * 8 + row_base + 2]), 
                    fabs(r[i * 8 + row_base + 3]), 
                    fabs(r[i * 8 + row_base + 4]), 
                    fabs(r[i * 8 + row_base + 5]), 
                    fabs(r[i * 8 + row_base + 6]), 
                    fabs(r[i * 8 + row_base + 7])
                );
                _local_residual += simd_reduc_plusd(simd_abs_r);
            }

            for (int i = (len / 8) * 8; i < len; i++) {
                _local_residual += fabs(r[i + row_base]);
            }    

        }
        double reduce_buf[64];
        athread_ssync_array();

        athread_redurt(&_local_residual, &residual, 1, athread_double, OP_add, reduce_buf, 64);
    }

    double const init_residual = residual;

    if (fabs(residual / param.normfactor) > param.tolerance) {
        do {
            if (iter == 0) {
                // z = M(-1) * r
                pcg_precondition_csr(r, z);
                // tol_0= swap(r) * z

                sumprod = pcg_gsumProd_imbalanced(r, z);
                // p = z，每个从核复制一部分到共享 LDM 中
                for (int i = row_base, end = matrix_slice.rows + row_base; i != end; ++i) {
                    p[i] = z[i];
                }
            } else {
                sumprod_old = sumprod;
                // z = M(-1) * r
                pcg_precondition_csr(r, z);
                // tol_0= swap(r) * z
                sumprod = pcg_gsumProd_imbalanced(r, z);
                // beta = tol_1 / tol_0
                beta = sumprod / sumprod_old;
                // p = z + beta * p
                for (int i = row_base, end = matrix_slice.rows + row_base; i != end; ++i) {
                    p[i] = z[i] + beta * p[i];
                }
            }

            athread_ssync_array();
            // Ax = A * p
            csr_spmv(p, Ax);
            // alpha = tol_0 / tol_1 = (swap(r) * z) / ( swap(p) * A * p)
            alpha = sumprod / pcg_gsumProd_imbalanced(p, Ax);

            // x = x + alpha * p
            // r = r - alpha * Ax
            double _local_residual = 0.0;
            for (int i = row_base, end = matrix_slice.rows + row_base; i != end; ++i) {
                x[i] = x[i] + alpha * p[i];
                r[i] = r[i] - alpha * Ax[i];
            }

            {
                int len = matrix_slice.rows;
                doublev8 simd_abs_r;
                for (int i = 0; i < len / 8; i++) {
                    simd_abs_r = simd_set_doublev8(
                        fabs(r[i * 8 + row_base]), 
                        fabs(r[i * 8 + row_base + 1]), 
                        fabs(r[i * 8 + row_base + 2]), 
                        fabs(r[i * 8 + row_base + 3]), 
                        fabs(r[i * 8 + row_base + 4]), 
                        fabs(r[i * 8 + row_base + 5]), 
                        fabs(r[i * 8 + row_base + 6]), 
                        fabs(r[i * 8 + row_base + 7])
                    );
                    _local_residual += simd_reduc_plusd(simd_abs_r);
                }

                for (int i = (len / 8) * 8; i < len; i++) {
                    _local_residual += fabs(r[i + row_base]);
                }    

            }

            double reduce_buf[64];
            athread_ssync_array();
            athread_redurt(&_local_residual, &residual, 1, athread_double, OP_add, reduce_buf, 64);
        } while (++iter < param.maxIter && (residual / param.normfactor) >= param.tolerance);
    }

    struct PCGReturn result;
    result.residual = residual;
    result.iter = iter;
    result.init_residual = init_residual;

    athread_dma_put(param.result, &result, sizeof(result));
    free_csr_matrix_in_ldm();
}
